import os

GPU_IDX = int(os.environ.get('CUDA_VISIBLE_DEVICES', 0))

import time
from PyCmpltrtok.common import sep
from PyCmpltrtok.common_torch import check_dtype
from xcommon import N_DATA

import numpy as np
from datasets import Dataset


seq_len, dataset_size = 512, N_DATA
dummy_data = {
    "input_ids": np.random.randint(100, 30000, (dataset_size, seq_len)),
    "labels": np.random.randint(0, 1, (dataset_size)),
}
ds = Dataset.from_dict(dummy_data)
ds.set_format("pt")

from pynvml import *


def print_gpu_utilization():
    nvmlInit()
    handle = nvmlDeviceGetHandleByIndex(GPU_IDX)
    info = nvmlDeviceGetMemoryInfo(handle)
    print(f"GPU memory occupied: {info.used//1024**2} MB.")


def print_summary(result):
    print(f"Time: {result.metrics['train_runtime']:.2f}")
    print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
    print_gpu_utilization()
    
    
print_gpu_utilization()

import torch


torch.ones((1000, 1000)).to("cuda")
print_gpu_utilization()

from transformers import AutoModelForSequenceClassification


model = AutoModelForSequenceClassification.from_pretrained("/home/yunpeng/models/hf/bert-large-uncased/6da4b6a26a1877e173fca3225479512db81a5e5b").to("cuda")
sep('sleep')
time.sleep(0.75)
print_gpu_utilization()
dtypes = check_dtype(model)
print(f'dtypes: {dtypes}')

default_args = {
    "output_dir": "ckpt.tmp.d",
    "evaluation_strategy": "steps",
    "num_train_epochs": 1,
    "log_level": "debug",
    "report_to": "none",
}

from transformers import TrainingArguments, Trainer, logging
import transformers

logging.set_verbosity_debug()

class MyCallback(transformers.TrainerCallback):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    """
    https://discuss.huggingface.co/t/logs-of-training-and-validation-loss/1974/3
    """

    def on_step_end(self, args, state, control, **kwargs):
        print_gpu_utilization()  

training_args = TrainingArguments(per_device_train_batch_size=4, **default_args)
trainer = Trainer(model=model, args=training_args, train_dataset=ds, callbacks=[MyCallback()])
result = trainer.train()
print_summary(result)
